package algorithm.middle;

import pojo.TreeNode;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class DistanceK863 {

    public List<Integer> res = new ArrayList<>();
    public Map<Integer, TreeNode> map = new HashMap<>();
    public int target;
    public TreeNode curNode;


    //[0,null,1,null,2,null,3]  1  2
    /*          0
     *              1
     *                  2
     *                      3
     */
    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        this.target = target.val;
        map.put(root.val,null);
        foreach(root);
        child(curNode,k);
        parent(curNode,k);
        return res;
    }

    public void parent(TreeNode root,int k){
        TreeNode parent = map.get(root.val);
        if (parent==null) return;
        k = k-1;
        if (k == 0) {
            res.add(parent.val);
            return;
        }
        parent(parent,k);
        k = k -1;
        if (parent.left!=null && parent.left.val!=root.val){
            child(parent.left,k);
        }else if (parent.right!=null && parent.right.val!=root.val){
            child(parent.right,k);
        }
    }

    public void child(TreeNode root,int k){
       if (k==0){
           res.add(root.val);
       }
       k = k -1;
       if (root.left!=null){
           child(root.left,k);
       }
       if (root.right!=null){
           child(root.right,k);
       }
    }

    public void foreach(TreeNode root) {
        if (root.val == target){
            curNode = root;
        }
        if (root.left!=null){
            map.put(root.left.val,root);
            foreach(root.left);
        }
        if (root.right!=null){
            map.put(root.right.val,root);
            foreach(root.right);
        }
    }
}
